from PIL import Image
import time
import logging

def text_detection(image_paths, batch_size=1, save=False):
    from paddlex import create_model
    logging.info("开始文本识别任务")
    logging.info("初始化模型...")
    t0 = time.time()
    model = create_model(model_name="PP-OCRv4_server_det")
    t1 = time.time()
    logging.info(f"模型初始化耗时: {t1 - t0:.2f} 秒")

    logging.info("模型推理...")
    t2 = time.time()
    output = model.predict(image_paths, batch_size=batch_size)
    t3 = time.time()
    logging.info(f"推理耗时: {t3 - t2:.2f} 秒")
    output = list(output)
    results = []
    for i, res in enumerate(output):
        result = {}
        if save:
            res.print()
            res.save_to_img(save_path="./output/")
            res.save_to_json(save_path="./output/res.json")
        result["dt_polys"] = res['dt_polys']
        result["dt_scores"] = res['dt_scores']
        result['img_size'] = (res['input_img'].shape[1], res['input_img'].shape[0])
        results.append(result)
    t4 = time.time()
    logging.info(f"结果解析耗时: {t4 - t3:.2f} 秒")
    logging.info(f"text_detection 总耗时: {t4 - t0:.2f} 秒")
    return results

if __name__ == "__main__":
    image_paths = [
        "/home/fengjie/doc-parser/MinerU/src/input_doc/sample/提取自德州德达城市建设投资运营有限公司2023年面向专业投资者公开发行公司债券（第一期）募集说明书(1).pdf"
    ]
    results = text_detection(image_paths,batch_size=1,save=True)
    print(results) 